﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Net;
using System.Net.Sockets;
using System.Reflection;
using System.Diagnostics;
using System.Threading;
using System.Net.NetworkInformation;

namespace DServiceDiscovery
{
	public class DServiceDiscovery: IDisposable
	{
		public const int DefaultPort = 6335;

		private UdpClient client;
		private object clientLock = new object();
		public int MaxPort { get; set; }
		public uint AliveScanDelay { get; set; }
		public uint FullScanDelay { get; set; }
		private IPEndPoint localEP;

		private Dictionary<string, List<IPEndPoint>> serviceEndpoints = new Dictionary<string,List<IPEndPoint>>();

		public DServiceDiscovery(IPEndPoint localEP)
		{
			this.localEP = localEP;
			SetDefaults();
		}

		private void SetDefaults()
		{
			AliveScanDelay = 60;
			FullScanDelay = 60 * 10;
		}

		public DServiceDiscovery(int port)
		{
			this.localEP = new IPEndPoint(IPAddress.Any, port);
			SetDefaults();
		}

		public DServiceDiscovery()
		{
			this.localEP = new IPEndPoint(IPAddress.Any, DefaultPort);
			SetDefaults();
		}

		private Thread scanThread;
		private void ScanThread()
		{
			var fullScan = new Stopwatch();
			fullScan.Start();
			var aliveScan = new Stopwatch();
			aliveScan.Start();

			this.serviceEndpoints = FullScan();

			while (true)
			{
				Thread.Sleep((int)Math.Min(1000, Math.Min(AliveScanDelay, FullScanDelay)));
				if (fullScan.ElapsedMilliseconds / 1000 > FullScanDelay)
				{
					fullScan.Reset();
					fullScan.Start();
					aliveScan.Reset();
					aliveScan.Start();
					this.serviceEndpoints = FullScan();
				}
				else if (aliveScan.ElapsedMilliseconds / 1000 > AliveScanDelay)
				{
					aliveScan.Reset();
					aliveScan.Start();
					// this.serviceEndpoints = ScanAlive();
				}
			}
		}
		private Dictionary<string, List<IPEndPoint>> ScanAlive()
		{
			var alive = this.serviceEndpoints;
			var result = new Dictionary<string, List<IPEndPoint>>();

			foreach (var service in alive)
			{
				foreach(var endpoint in service.Value)
				{
					AddEndpointServices(result, endpoint);
				}
				if (Debugger.IsAttached)
					Debug.WriteLine("alive scan: " + service.Key);
			}

			return result;
		}

		private void AddEndpointServices(Dictionary<string, List<IPEndPoint>> result, IPEndPoint endpoint)
		{
			var services = Scan(endpoint);
			lock (result)
			{
				foreach (var serviceName in services)
				{
					List<IPEndPoint> endpoints;
					if (result.TryGetValue(serviceName, out endpoints))
					{
						endpoints.Add(endpoint);
					}
					else
					{
						endpoints = new List<IPEndPoint>();
						result[serviceName] = endpoints;
					}
				}
			}
		}

		private static IPAddress[] GetHostIPs()
		{
			var hostName = Dns.GetHostName();
			var entry = Dns.GetHostEntry(hostName);
			return entry.AddressList;
		}

		private Dictionary<string, List<IPEndPoint>> FullScan()
		{
			var result = new Dictionary<string, List<IPEndPoint>>();
			var start = new Stopwatch();
			start.Start();

			var ipProperties = IPGlobalProperties.GetIPGlobalProperties();
			var tcp = ipProperties.GetActiveTcpListeners();
			var udp = ipProperties.GetActiveUdpListeners();
			var distinct = tcp.Concat(udp);
			var ips = GetHostIPs();
			var endpoints = distinct
				.SelectMany(endpoint =>
					{
						if (endpoint.Address.Equals(IPAddress.Any) || endpoint.Address.Equals(IPAddress.IPv6Any))
						{
							return ips
								.Select(ip => new IPEndPoint(ip, endpoint.Port));
						}
						else
							return Enumerable.Range(0, 1)
								.Select(i => endpoint);
					}).Distinct();

			Debug.WriteLineIf(Debugger.IsAttached, "endpoints were computed");

			var pairs = endpoints.AsParallel(16)
				.SelectMany(
					ep => Scan(ep)
						.Select(service => new { Service = service, EndPoint = ep})
				);

			Debug.WriteLineIf(Debugger.IsAttached, "pairs were found");

			foreach (var service in pairs)
			{
				List<IPEndPoint> eps;
				if (result.TryGetValue(service.Service, out eps))
				{
					if (!eps.Contains(service.EndPoint))
						eps.Add(service.EndPoint);
				}else
					result[service.Service] = new List<IPEndPoint>{ service.EndPoint};
			}

			//GetHostIPs().AsParallel().ForAll(ip =>
			//{
			//    Enumerable.Range(IPEndPoint.MinPort, MaxPort - IPEndPoint.MinPort + 1).AsParallel().ForAll(port =>
			//    {
			//        AddEndpointServices(result, new IPEndPoint(ip, port));
			//    });
			//});

			//if (Debugger.IsAttached)
			//    Debug.WriteLine("done fullscan");
			Debug.WriteLineIf(Debugger.IsAttached, "full scan complete in " + start.ElapsedMilliseconds / 1000 + "sec");

			return result;
		}

		static readonly Dictionary<int, string>
			tcpPorts = new Dictionary<int, string>(),
			udpPorts = new Dictionary<int, string>();

		static DServiceDiscovery()
		{
			tcpPorts[8]		= "ICMP";
			tcpPorts[21]	= "FTP";
			tcpPorts[22]	= "SSH";
			tcpPorts[23]	= "Telnet";
			tcpPorts[25]	= "SMTP";
			tcpPorts[37]	= "TIME";
			tcpPorts[43]	= "WhoIs";
			tcpPorts[53]	= "DNS";
			tcpPorts[80]	= "HTTP";
			tcpPorts[110]	= "POP3";
			tcpPorts[1080]	= "Socks";
			tcpPorts[1194]	= "OpenVPN";

			udpPorts[8]		= "ICMP";
			udpPorts[21]	= "FTP";
			udpPorts[22]	= "SSH";
			udpPorts[23]	= "Telnet";
			udpPorts[25]	= "SMTP";
			udpPorts[37]	= "TIME";
			udpPorts[43]	= "WhoIs";
			udpPorts[53]	= "DNS";
			udpPorts[80]	= "HTTP";
			udpPorts[110]	= "POP3";
			udpPorts[1080]	= "Socks";
			udpPorts[1194]	= "OpenVPN";
			udpPorts[DefaultPort] = "DServiceDiscovery";
		}

		private string[] Scan(IPEndPoint endpoint)
		{
			// This is primitive scan stub based on port numbers
			List<string> result = new List<string>();
			try
			{
				using(var scannerTcpClient = new TcpClient(endpoint.AddressFamily))
				{
					scannerTcpClient.Connect(endpoint);
					scannerTcpClient.GetStream().Close();
				}

				string service;
				if (tcpPorts.TryGetValue(endpoint.Port, out service))
					result.Add(service);
				else
					result.Add("unknown");
			}
			catch (SocketException)
			{
			}

			try
			{
				bool available;
				using (var scannerUdpClient = new UdpClient(endpoint.AddressFamily))
				{
					scannerUdpClient.Send(new byte[] { }, 0, endpoint);
					var stopwatch = new Stopwatch();
					stopwatch.Start();
					while (scannerUdpClient.Available == 0 && stopwatch.ElapsedMilliseconds < 5)
						Thread.Sleep(0);

					available = scannerUdpClient.Available > 0;
				}

				if (available)
				{
					string service;
					if (udpPorts.TryGetValue(endpoint.Port, out service))
						result.Add(service);
					else
						result.Add("unknown");
				}
			}
			catch (SocketException) { }

			if (Debugger.IsAttached)
				Debug.WriteLine("scanned: " + endpoint);

			return result.ToArray();
		}

		public void Run()
		{
			lock (clientLock)
			{
				if (client != null)
					throw new InvalidOperationException("already running");

				if (localEP == null)
					throw new ObjectDisposedException("DServiceDiscovery");

				client = new UdpClient(localEP);
				localEP = null;
			}

			scanThread = new Thread(new ThreadStart(ScanThread));
			scanThread.Start();

			try
			{
				while (client != null)
				{
					IPEndPoint requestSource = null;

					byte[] data;
					try
					{
						data = client.Receive(ref requestSource);
					}
					catch (SocketException)
					{
						lock (clientLock)
						{
							client.Close();
							client = null;
						}
						break;
					}

					try
					{
						var request = Encoding.UTF8.GetString(data);
						HandleRequest(request, requestSource);
					}
					catch (Exception)
					{
						client.Send(new byte[] { }, 0, requestSource);
					}
				}
			}
			catch (NullReferenceException)
			{
				if (client != null) throw;
			}
			catch (ObjectDisposedException)
			{
				client = null;
			}

			while (scanThread.ThreadState != System.Threading.ThreadState.Stopped)
			{
				scanThread.Abort();
				Thread.Sleep(50);
			}
			scanThread = null;
		}

		public void Stop()
		{
			client.Close();
			client = null;
			while (scanThread != null)
				Thread.Sleep(0);
		}

		public static class Commands
		{
			public const string Version = "ver";
			public const string GetService = "get-service";
			public const string GetServices = "get-services";
		}

		private void HandleRequest(string request, IPEndPoint requestSource)
		{
			string response;

			if (request.StartsWith(Commands.Version))
				response = GetVersionString();
			else if (request.StartsWith(Commands.GetService))
				response = GetService(request.Substring(Commands.GetService.Length + 1));
			else if (request.StartsWith(Commands.GetServices))
				response = GetServices();
			else
				response = "unknown command";

			if (response == null)
				return;

			byte[] data = Encoding.UTF8.GetBytes(response);
			client.Send(data, data.Length, requestSource);
		}

		internal string GetServices()
		{
			var result = new StringBuilder();
			var endpoints = this.serviceEndpoints;
			foreach (var service in endpoints)
			{
				result.AppendLine(service.Key);
				foreach (var endpoint in service.Value)
					result.AppendLine(endpoint.ToString());
				result.AppendLine();
			}
			return result.ToString();
		}

		private string GetService(string serviceName)
		{
			var endpoints = this.serviceEndpoints;
			List<IPEndPoint> results;
			if (endpoints.TryGetValue(serviceName, out results))
			{
				StringBuilder response = new StringBuilder();
				foreach(var result in results)
					response.AppendLine(result.ToString());

				return response.ToString();
			} else
				return null;
		}

		private string GetVersionString()
		{
			var thisAssembly = Assembly.GetExecutingAssembly();
			return thisAssembly.GetName().Version.ToString();
		}

		#region IDisposable Members

		public void Dispose()
		{
			Stop();

			lock (clientLock)
				if (this.client != null)
					this.client.Close();
		}

		~DServiceDiscovery()
		{
			Dispose();
		}
		#endregion
	}
}
